import argparse
import json
import os
import sys
import time
from pathlib import Path

import cv2
import numpy as np

try:
    import pyrealsense2 as rs
except ImportError as exc:
    print("pyrealsense2 未安装。请先安装 Intel RealSense SDK Python 包: pip install pyrealsense2", file=sys.stderr)
    raise


def ensure_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def write_orbslam3_yaml_from_intrinsics(yaml_path: Path,
                                         width: int,
                                         height: int,
                                         fps: int,
                                         fx: float,
                                         fy: float,
                                         cx: float,
                                         cy: float,
                                         dist_coeffs: np.ndarray,
                                         depth_map_factor: float = 1000.0) -> None:
    # ORB-SLAM3 RGB-D 配置（最小必需项 + 常用 ORB 参数），相机畸变按 k1,k2,p1,p2,k3 顺序
    k1, k2, p1, p2, k3 = 0.0, 0.0, 0.0, 0.0, 0.0
    if dist_coeffs is not None and len(dist_coeffs) >= 5:
        k1, k2, p1, p2, k3 = [float(x) for x in dist_coeffs[:5]]

    content = f"""
# Auto-generated for RealSense RGB-D recording
Camera.type: "PinHole"

Camera.width: {width}
Camera.height: {height}

Camera.fx: {fx:.8f}
Camera.fy: {fy:.8f}
Camera.cx: {cx:.8f}
Camera.cy: {cy:.8f}

Camera.k1: {k1:.8f}
Camera.k2: {k2:.8f}
Camera.p1: {p1:.8f}
Camera.p2: {p2:.8f}
Camera.k3: {k3:.8f}

Camera.fps: {fps}
Camera.RGB: 1

# 对于以 16-bit PNG 保存的毫米深度图，DepthMapFactor 设为 1000.0
DepthMapFactor: {depth_map_factor:.1f}

# ORB 提取器参数（可按需要微调）
ORBextractor.nFeatures: 1200
ORBextractor.scaleFactor: 1.2
ORBextractor.nLevels: 8
ORBextractor.iniThFAST: 20
ORBextractor.minThFAST: 7

# Viewer（可选）
Viewer.KeyFrameSize: 0.05
Viewer.KeyFrameLineWidth: 1
Viewer.GraphLineWidth: 1
Viewer.PointSize: 2
Viewer.CameraSize: 0.08
Viewer.CameraLineWidth: 3
Viewer.ViewpointX: 0
Viewer.ViewpointY: -0.7
Viewer.ViewpointZ: -1.8
Viewer.ViewpointF: 500
"""
    yaml_path.write_text(content.strip() + "\n", encoding="utf-8")


def record_realsense(output_dir: Path,
                     width: int,
                     height: int,
                     fps: int,
                     duration: float | None,
                     serial: str | None) -> None:
    ensure_dir(output_dir)
    rgb_dir = output_dir / "rgb"
    depth_dir = output_dir / "depth"
    ensure_dir(rgb_dir)
    ensure_dir(depth_dir)

    rgb_list = open(output_dir / "rgb.txt", "w", encoding="utf-8")
    depth_list = open(output_dir / "depth.txt", "w", encoding="utf-8")
    assoc_list = open(output_dir / "associations.txt", "w", encoding="utf-8")

    pipeline = rs.pipeline()
    cfg = rs.config()
    if serial:
        cfg.enable_device(serial)
    cfg.enable_stream(rs.stream.color, width, height, rs.format.bgr8, fps)
    cfg.enable_stream(rs.stream.depth, width, height, rs.format.z16, fps)

    profile = pipeline.start(cfg)
    align = rs.align(rs.stream.color)

    depth_sensor = profile.get_device().first_depth_sensor()
    depth_scale = float(depth_sensor.get_depth_scale())  # meters per unit

    color_stream = profile.get_stream(rs.stream.color).as_video_stream_profile()
    intr = color_stream.get_intrinsics()
    fx, fy = float(intr.fx), float(intr.fy)
    cx, cy = float(intr.ppx), float(intr.ppy)
    dist_coeffs = np.array(intr.coeffs, dtype=np.float64) if intr.coeffs else np.zeros(5)

    # 保存相机内参（JSON）
    intr_json = {
        "width": width,
        "height": height,
        "fps": fps,
        "fx": fx,
        "fy": fy,
        "cx": cx,
        "cy": cy,
        "distortion_model": str(intr.model),
        "distortion_coeffs": dist_coeffs.tolist(),
        "depth_scale_m": depth_scale,
        "depth_unit": "millimeter",
        "depth_map_factor": 1000.0
    }
    (output_dir / "camera_intrinsics.json").write_text(json.dumps(intr_json, indent=2), encoding="utf-8")

    # 生成 ORB-SLAM3 YAML（自动）
    write_orbslam3_yaml_from_intrinsics(output_dir / "orbslam3_realsense_rgbd.yaml",
                                         width, height, fps, fx, fy, cx, cy, dist_coeffs,
                                         depth_map_factor=1000.0)

    print("开始录制。窗口激活后按 Q 结束。")
    start_t = time.time()
    frame_idx = 0

    try:
        while True:
            frames = pipeline.wait_for_frames()
            aligned = align.process(frames)
            depth_frame = aligned.get_depth_frame()
            color_frame = aligned.get_color_frame()
            if not depth_frame or not color_frame:
                continue

            # 使用相机时间戳（毫秒）转换为秒，保证 RGB/Depth 一致
            ts_sec = float(color_frame.get_timestamp()) / 1000.0

            color = np.asanyarray(color_frame.get_data())  # HxWx3, BGR8
            depth = np.asanyarray(depth_frame.get_data()).astype(np.float32)  # uint16 raw units

            # 转为毫米并保存为 16-bit PNG
            depth_m = depth * depth_scale
            depth_mm = (depth_m * 1000.0)
            depth_mm = np.clip(depth_mm, 0, np.iinfo(np.uint16).max).astype(np.uint16)

            # 保存文件
            ts_str = f"{ts_sec:.6f}"
            rgb_filename = f"{ts_str}.png"
            depth_filename = f"{ts_str}.png"

            cv2.imwrite(str(rgb_dir / rgb_filename), color)
            cv2.imwrite(str(depth_dir / depth_filename), depth_mm)

            rgb_rel = f"rgb/{rgb_filename}"
            depth_rel = f"depth/{depth_filename}"
            rgb_list.write(f"{ts_str} {rgb_rel}\n")
            depth_list.write(f"{ts_str} {depth_rel}\n")
            assoc_list.write(f"{ts_str} {rgb_rel} {ts_str} {depth_rel}\n")

            frame_idx += 1

            # 简单预览（可选）
            preview = color.copy()
            cv2.putText(preview, f"{frame_idx}  {ts_str}s", (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            cv2.imshow("RealSense RGB-D Recording (Q to stop)", preview)
            key = cv2.waitKey(1) & 0xFF
            if key in (ord('q'), ord('Q')):
                break

            if duration is not None and (time.time() - start_t) >= duration:
                break
    finally:
        pipeline.stop()
        rgb_list.close()
        depth_list.close()
        assoc_list.close()
        cv2.destroyAllWindows()

    print(f"录制完成，共保存 {frame_idx} 帧。数据集位于: {output_dir}")
    print(f"已生成: {output_dir / 'orbslam3_realsense_rgbd.yaml'} 用于 ORB-SLAM3")


def main() -> None:
    parser = argparse.ArgumentParser(description="录制 RealSense RGB-D 为 TUM 数据集，并生成 ORB-SLAM3 配置")
    parser.add_argument("--out", required=True, help="输出数据集目录，例如 data/room1")
    parser.add_argument("--width", type=int, default=640)
    parser.add_argument("--height", type=int, default=480)
    parser.add_argument("--fps", type=int, default=30)
    parser.add_argument("--duration", type=float, default=None, help="录制时长（秒），为空则按 Q 手动停止")
    parser.add_argument("--serial", type=str, default=None, help="指定 RealSense 序列号，可选")

    args = parser.parse_args()

    out_dir = Path(args.out)
    record_realsense(out_dir, args.width, args.height, args.fps, args.duration, args.serial)


if __name__ == "__main__":
    main()


